package ex3.visitors;

import common.exceptions.SemanticError;
import ic.ast.Visitor;
import ic.ast.decl.*;
import ic.ast.expr.*;
import ic.ast.expr.BinaryOp.BinaryOps;
import ic.ast.expr.UnaryOp.UnaryOps;
import ic.ast.stmt.*;
import ic.symbols.*;
import ic.types.MethodType;
import ic.types.TypeTable;

import java.util.Iterator;

/**
 * Created with IntelliJ IDEA.
 * User: Roi
 * Date: 12/6/13
 * <p/>
 * A visitor responsible for type and scope checks
 */
public class TypeCheckVisitor implements Visitor {

    private boolean inVirtualMethodContext = false;

    @Override
    public Object visit(Program program) throws Exception {
        for (DeclClass icClass : program.getClasses()) {
            icClass.accept(this);
        }
        return null;
    }

    @Override
    public Object visit(DeclClass icClass) throws Exception {
        for (DeclField field : icClass.getFields()) {
            field.accept(this);
        }
        for (DeclMethod method : icClass.getMethods()) {
            method.accept(this);
        }
        return null;
    }

    @Override
    public Object visit(DeclField field) throws Exception {
        return TypeTable.getType(field.getType(), false);
    }

    @Override
    public Object visit(DeclVirtualMethod method) throws Exception {
        inVirtualMethodContext = true;
        for (Statement stmt : method.getStatements()) {
            stmt.accept(this);
        }
        inVirtualMethodContext = false;
        return null;
    }

    @Override
    public Object visit(DeclStaticMethod method) throws Exception {
        for (Statement stmt : method.getStatements()) {
            stmt.accept(this);
        }
        return null;
    }

    @Override
    public Object visit(DeclLibraryMethod method) throws Exception {
        return TypeTable.getType(method);
    }

    @Override
    public Object visit(Parameter formal) throws Exception {
        return TypeTable.getType(formal.getType(), false);
    }

    @Override
    public Object visit(PrimitiveType type) throws Exception {
        return TypeTable.getType(type, false);
    }

    @Override
    public Object visit(ClassType type) throws Exception {
        try {
            TypeTable.getClassTypeByName(type.getDisplayName());
        } catch (SemanticError semantic) {
            throw new SemanticError(type.getLine(), semantic.getMessage());
        }
        return TypeTable.getType(type);
    }

    /**
     * Check:
     * Assignment type is a subtype of the variable assigned to.
     * Returns null, or an error if found.
     */
    @Override
    public Object visit(StmtAssignment assignment) throws Exception {

        Type variableType = (Type) assignment.getVariable().accept(this);
        if (variableType == null)
            return null;

        Type assignmentType = (Type) assignment.getAssignment().accept(this);
        if (assignmentType == null)
            return null;

        if (!(assignmentType.subTypeOf(variableType))) {
            throw new SemanticError(assignment.getLine(), "Invalid assignment of type " + assignmentType.getOutputName()
                    + " to variable of type " + variableType.getOutputName());

        }
        return null;
    }

    @Override
    public Object visit(StmtCall callStatement) throws Exception {
        callStatement.getCall().accept(this);
        return null;
    }

    /**
     * Check: return value code is a subtype of the method type
     */
    @Override
    public Object visit(StmtReturn returnStatement) throws Exception {
        Type returnValueType;

        if (returnStatement.hasValue()) {
            returnValueType = (Type) returnStatement.getValue().accept(this);
        } else {
            returnValueType = TypeTable.getType("void");
        }

        String methodName = returnStatement.getEnclosingScope().getCurrentMethodName();
        MethodSymbol methodSymbol = (MethodSymbol) returnStatement.getEnclosingScope().lookup(methodName);
        MethodType methodType = methodSymbol.getMethodType();
        Type methodReturnType = TypeTable.getType(methodType.getReturnType());

        if (!returnValueType.subTypeOf(methodReturnType)) {
            throw new SemanticError(returnStatement.getLine(), "Return statement is not of type " + methodReturnType.getOutputName());
        }

        return null;
    }

    /**
     * Check: types in an If/Else statement.
     * Makes sure that the if condition is of type boolean.
     * Returns null, or an error if found.
     */
    @Override
    public Object visit(StmtIf ifStatement) throws Exception {
        Type conditionType = (Type) ifStatement.getCondition().accept(this);
        if (conditionType != TypeTable.getType("boolean")) {
            throw new SemanticError(ifStatement.getLine(), "Non boolean condition for if statement");
        }

        ifStatement.getOperation().accept(this);

        if (ifStatement.hasElse()) {
            ifStatement.getElseOperation().accept(this);
        }

        return null;
    }

    /**
     * Check: while statement condition is of type boolean.
     * Returns null, or an error if found.
     */
    @Override
    public Object visit(StmtWhile whileStatement) throws Exception {
        Type conditionType = (Type) whileStatement.getCondition().accept(this);
        if (conditionType != TypeTable.getType("boolean")) {
            throw new SemanticError(whileStatement.getLine(), "Non boolean condition for while statement");
        }

        whileStatement.getOperation().accept(this);
        return null;
    }

    @Override
    public Object visit(StmtBreak breakStatement) {
        return null;
    }

    @Override
    public Object visit(StmtContinue continueStatement) {
        return null;
    }

    @Override
    public Object visit(StmtBlock statementsBlock) throws Exception {
        for (Statement stmt : statementsBlock.getStatements()) {
            stmt.accept(this);
        }
        return null;
    }

    /**
     * Check: variable type is a subtype of the local variables type
     */
    @Override
    public Object visit(LocalVariable localVariable) throws Exception {
        localVariable.getType().accept(this);
        Type varType = TypeTable.getType(localVariable.getType(), false);

        if (localVariable.getInitialValue() != null) {
            Type initValueType = (Type) localVariable.getInitialValue().accept(this);
            if (!initValueType.subTypeOf(varType)) {
                throw new SemanticError(localVariable.getLine(), "type mismatch: expected " + varType + " and got " + initValueType);
            }
        }

        return null;
    }

    /**
     * Return location type, or an error if found.
     */
    @Override
    public Object visit(RefVariable location) throws Exception {
        Type expressionType;

        if (TypeTable.getType(location.getName()) == null) {
            Symbol varSymbol = location.getEnclosingScope().lookup(location.getName());
            if (varSymbol == null) {
                throw new SemanticError(location.getLine(), location.getName() + " not found in symbol table");
            }
            if (varSymbol.getSymbolType() == SymbolType.VARIABLE) {
                if (varSymbol.getLine() > location.getLine() ||
                        (varSymbol.getLine() == location.getLine())) {
                    Symbol altSymbol = location.getEnclosingScope().getParent().lookup(location.getName());
                    if (altSymbol != null && (altSymbol.getSymbolType() == SymbolType.VARIABLE || altSymbol.getSymbolType() == SymbolType.FIELD)) {
                        varSymbol = altSymbol;
                    } else {
                        throw new SemanticError(location.getLine(), "variable " + location.getName() + " is used before declaration");
                    }
                }
            }
            if (varSymbol.getSymbolType() == SymbolType.VARIABLE) {
                expressionType = TypeTable.getType(((VarSymbol) varSymbol).getType(), false);
            } else if (varSymbol.getSymbolType() == SymbolType.FIELD && inVirtualMethodContext) {
                expressionType = TypeTable.getType(((FieldSymbol) varSymbol).getType(), false);
            } else {
                throw new SemanticError(location.getLine(), "Use of field inside static method is not allowed");
            }
        } else {
            Type classType = TypeTable.getType(location.getName());
            Symbol fieldSymbol;
            try {
                SymbolTable classSymbolTable = TypeTable.getClassTypeByName(classType.getDisplayName()).getEnclosingScope().getSymbolTable(classType.getDisplayName());
                fieldSymbol = classSymbolTable.lookup(location.getName());
            } catch (SemanticError semantic) {
                throw new SemanticError(location.getLine(), "no such class: " + classType.getDisplayName());
            }
            if (fieldSymbol == null) {
                throw new SemanticError(location.getLine(), location.getName() + " not found in symbol table");
            }
            if (fieldSymbol.getSymbolType() != SymbolType.FIELD) {
                throw new SemanticError(location.getLine(), "illegal reference: " + location.getName());
            }
            expressionType = TypeTable.getType(((FieldSymbol) fieldSymbol).getType(), false);
        }

        return expressionType;
    }

    @Override
    public Object visit(RefField location) throws Exception {
        Type expressionType;

        if (location.getObject() == null) {
            Symbol varSymbol = location.getEnclosingScope().lookup(location.getField());
            if (varSymbol == null) {
                throw new SemanticError(location.getLine(), location.getField() + " not found in symbol table");
            }
            if (varSymbol.getSymbolType() == SymbolType.VARIABLE) {
                if (varSymbol.getLine() > location.getLine() ||
                        (varSymbol.getLine() == location.getLine())) {
                    Symbol altSymbol = location.getEnclosingScope().getParent().lookup(location.getField());
                    if (altSymbol != null && (altSymbol.getSymbolType() == SymbolType.VARIABLE || altSymbol.getSymbolType() == SymbolType.FIELD)) {
                        varSymbol = altSymbol;
                    } else {
                        throw new SemanticError(location.getLine(), "variable " + location.getField() + " is used before declaration");
                    }
                }
            }
            if (varSymbol.getSymbolType() == SymbolType.VARIABLE) {
                expressionType = TypeTable.getType(((VarSymbol) varSymbol).getType(), false);
            } else if (varSymbol.getSymbolType() == SymbolType.FIELD && inVirtualMethodContext) {
                expressionType = TypeTable.getType(((FieldSymbol) varSymbol).getType(), false);
            } else {
                throw new SemanticError(location.getLine(), "illegal reference: " + location.getField());
            }
        } else {
            Type classType = (Type) location.getObject().accept(this);
            Symbol fieldSymbol;
            try {
                SymbolTable classSymbolTable = TypeTable.getClassTypeByName(classType.getDisplayName()).getEnclosingScope().getSymbolTable(classType.getDisplayName());
                fieldSymbol = classSymbolTable.lookup(location.getField());
            } catch (SemanticError semantic) {
                throw new SemanticError(location.getLine(), "no such class: " + classType.getDisplayName());
            }
            if (fieldSymbol == null) {
                throw new SemanticError(location.getLine(), location.getField() + " not found in symbol table");
            }
            if (fieldSymbol.getSymbolType() != SymbolType.FIELD) {
                throw new SemanticError(location.getLine(), "illegal reference: " + location.getField());
            }
            expressionType = TypeTable.getType(((FieldSymbol) fieldSymbol).getType(), false);
        }

        return expressionType;
    }

    /**
     * Check: array index is of type int.
     * Return element type, or an error if found.
     */
    @Override
    public Object visit(RefArrayElement location) throws Exception {
        Type arrayType = (Type) location.getArray().accept(this);
        Type indexType = (Type) location.getIndex().accept(this);

        if (indexType != TypeTable.getType("int")) {
            throw new SemanticError(location.getLine(), "array index must be of type int");
        }

        if (arrayType.getArrayDimension() < 1) {
            throw new SemanticError(location.getLine(), arrayType + " is not an array type");
        }

        String typeName = arrayType.toString();
        String elementType = typeName.substring(0, typeName.length() - 2);

        return TypeTable.getType(elementType);

    }

    /**
     * Check:
     * types of arguments
     * method is defined in the correct class
     * all arguments correspond to the method's arguments types
     * Return method return value, or an error if found.
     */
    @Override
    public Object visit(StaticCall call) throws Exception {
        String className = call.getClassName();
        DeclClass classInstance;

        try {
            classInstance = TypeTable.getClassTypeByName(className);
        } catch (SemanticError e) {
            if (className.equals("Library")) {
                throw new SemanticError(call.getLine(), "Method " + call.getMethod() + " doesn't exist");
            } else {
                throw new SemanticError(call.getLine(), className + ": no such class");
            }
        }

        SymbolTable classScope = classInstance.getEnclosingScope().getSymbolTable(className);
        String staticMethodName = call.getMethod();

        Symbol methodSymbol = classScope.lookup(staticMethodName);

        if (methodSymbol == null || methodSymbol.getSymbolType() != SymbolType.METHOD ||
                !((MethodSymbol) methodSymbol).isStatic()) {
            throw new SemanticError(call.getLine(), "Method " + staticMethodName + " doesn't exist");
        }

        MethodType methodType = ((MethodSymbol) methodSymbol).getMethodType();
        if (call.getArguments().size() != methodType.getParamTypes().size()) {
            throw new SemanticError(call.getLine(), "Invalid number of arguments for " + className + "." + staticMethodName);
        }

        Iterator<Type> methodParams = methodType.getParamTypes().iterator();
        for (Expression arg : call.getArguments()) {
            Type paramType = TypeTable.getType(methodParams.next(), false);
            Type argType = (Type) arg.accept(this);
            if (!argType.subTypeOf(paramType)) {
                throw new SemanticError(call.getLine(), "argument and parameter type mismatch in call to " + staticMethodName +
                        ". Expected " + paramType + " and got " + argType);
            }
        }

        return TypeTable.getType(methodType.getReturnType(), false);
    }

    /**
     * Check:
     * method is defined in the relevant class
     * all arguments correspond to the method's arguments types
     * Return method return value, or an error if found.
     */
    @Override
    public Object visit(VirtualCall call) throws Exception {
        DeclClass instanceClass;
        Type instanceType;
        SymbolTable classScope;
        String instanceClassName;
        Expression location = call.getObject();

        if (call.hasExplicitObject() && !(location instanceof This)) {
            instanceType = (Type) location.accept(this);
            instanceClassName = instanceType.getDisplayName();
            try {
                instanceClass = TypeTable.getClassTypeByName(instanceClassName);
            } catch (SemanticError e) {
                throw new SemanticError(call.getLine(), e.getMessage());
            }
            classScope = instanceClass.getEnclosingScope().getSymbolTable(instanceClassName);
        } else {
            try {
                classScope = call.getEnclosingScope();
                if (classScope instanceof BlockSymbolTable) {
                    classScope = classScope.getParent();
                }
                if (classScope instanceof MethodSymbolTable) {
                    classScope = classScope.getParent();
                }
                instanceClass = TypeTable.getClassTypeByName(classScope.getName());
                instanceClassName = instanceClass.getName();
            } catch (SemanticError e) {
                throw new SemanticError(call.getLine(), e.getMessage());
            }
        }

        String methodName = call.getMethod();
        Symbol methodSymbol = classScope.lookup(methodName);

        if (methodSymbol == null || methodSymbol.getSymbolType() != SymbolType.METHOD) {
            throw new SemanticError(call.getLine(), methodName + " not found in symbol table");
        }
        if (!call.hasExplicitObject() && !inVirtualMethodContext && !((MethodSymbol) methodSymbol).isStatic()) {
            throw new SemanticError(call.getLine(), "Calling a local virtual method from inside a static method is not allowed");
        }
        MethodType methodType = ((MethodSymbol) methodSymbol).getMethodType();
        if (call.getArguments().size() != methodType.getParamTypes().size()) {
            throw new SemanticError(call.getLine(), "invalid amount of arguments passed to " + methodName);
        }

        Iterator<Type> methodParams = methodType.getParamTypes().iterator();
        for (Expression arg : call.getArguments()) {
            Type paramType = TypeTable.getType(methodParams.next(), false);
            Type argType = (Type) arg.accept(this);
            if (!argType.subTypeOf(paramType)) {
                throw new SemanticError(call.getLine(), "Method " + instanceClassName + "." + methodName +
                        " is not applicable for the arguments given");
            }
        }

        return TypeTable.getType(methodType.getReturnType(), false);
    }

    @Override
    public Object visit(This thisExpression) throws Exception {
        return TypeTable.getType(new ClassType(thisExpression.getLine(), thisExpression.getEnclosingScope().getName()));
    }

    /**
     * Check: new class expression - class type is legal and exists.
     */
    @Override
    public Object visit(NewInstance newClass) throws Exception {
        String className = newClass.getName();
        return TypeTable.getType(new ClassType(newClass.getLine(), className));
    }

    /**
     * Check: new array expression - element type is a legal type and that the size is of type int.
     * Return the array type.
     */
    @Override
    public Object visit(NewArray newArray) throws Exception {
        newArray.getType().accept(this);
        Type tmpType = TypeTable.getType(newArray.getType(), false);
        Type arrayType = TypeTable.getType(tmpType.clone(tmpType.getArrayDimension() + 1));

        Type sizeType = (Type) newArray.getSize().accept(this);
        if (sizeType != TypeTable.getType("int")) {
            throw new SemanticError(newArray.getLine(), "array size must be of type int");
        }

        return arrayType;
    }

    @Override
    public Object visit(Length length) throws Exception {
        Type arrType = (Type) length.getArray().accept(this);

        if (arrType.getArrayDimension() < 1) {
            throw new SemanticError(length.getLine(), "length is applicable only for array types");
        }

        return TypeTable.getType("int");
    }

    /**
     * Check: types in math binary operations.
     * for + operation: gets either 2 ints or 2 strings. Return int or string.
     * for {/,*,-} operations, gets 2 ints and return int.
     * Send an error otherwise
     */
    @Override
    public Object visit(BinaryOp binaryOp) throws Exception {
        Type lhsType = (Type) binaryOp.getFirstOperand().accept(this);
        Type rhsType = (Type) binaryOp.getSecondOperand().accept(this);
        BinaryOps operator = binaryOp.getOperator();

        // If operator is logical
        if (operator == BinaryOps.LAND || operator == BinaryOps.LOR
                || operator == BinaryOps.GT || operator == BinaryOps.GTE
                || operator == BinaryOps.LT || operator == BinaryOps.LTE
                || operator == BinaryOps.EQUAL || operator == BinaryOps.NEQUAL) {


            if ((operator == BinaryOps.LAND) || (operator == BinaryOps.LOR)) {
                if (lhsType != TypeTable.getType("boolean") || rhsType != TypeTable.getType("boolean")) {
                    throw new SemanticError(binaryOp.getLine(), "Logical " + operator.getDescription() +
                            " is applicable only for boolean operands");
                }
            }

            if (operator == BinaryOps.EQUAL || operator == BinaryOps.NEQUAL) {
                if (!(lhsType.subTypeOf(rhsType)) && !(rhsType.subTypeOf(lhsType))) {
                    throw new SemanticError(binaryOp.getLine(),
                            "comparison of foreign types with the operator: " + operator.getDescription());
                }
            }

            if (operator == BinaryOps.LTE || operator == BinaryOps.LT
                    || operator == BinaryOps.GTE || operator == BinaryOps.GT) {
                if (lhsType != TypeTable.getType("int") || rhsType != TypeTable.getType("int")) {
                    throw new SemanticError(binaryOp.getLine(), "Invalid logical binary op (" + operator + ") on non-integer expression");
                }
            }
            return TypeTable.getType("boolean");
        }

        if ((lhsType == null) || (rhsType == null)) return null;

        if ((lhsType != TypeTable.getType("null")) || (rhsType != TypeTable.getType("null"))) {
            if (rhsType == TypeTable.getType("null"))
                rhsType = TypeTable.getType("string");

            if (lhsType == TypeTable.getType("null"))
                lhsType = TypeTable.getType("string");
        }

        if (lhsType != rhsType)
            throw new SemanticError(binaryOp.getLine(), "Illegal " + binaryOp.getOperator().toString() +
                    " operation. Both operands must be of the same type " +
                    (binaryOp.getOperator() == BinaryOp.BinaryOps.PLUS ? "(int or string)" : "(int)"));

        if (binaryOp.getOperator() != BinaryOp.BinaryOps.PLUS) {
            if (lhsType != TypeTable.getType("int")) {
                throw new SemanticError(binaryOp.getLine(), binaryOp.getOperator().toString() + " operations are only for two operands of type int");
            }
        } else {
            if (lhsType != TypeTable.getType("int") && lhsType != TypeTable.getType("string"))
                throw new SemanticError(binaryOp.getLine(), binaryOp.getOperator().toString() + " operation is only for int or string operands");
        }

        return TypeTable.getType(lhsType, false);
    }

    /**
     * Return the type of the literal
     */
    @Override
    public Object visit(Literal literal) {
        switch (literal.getType()) {
            case STRING:
                return TypeTable.getType("string");
            case INT:
                return TypeTable.getType("int");
            case BOOLEAN:
                return TypeTable.getType("boolean");
            case VOID:
                return TypeTable.getType("null");
        }
        return null;
    }

    public Object visit(UnaryOp unaryOp) throws Exception {
        Type type = (Type) unaryOp.getOperand().accept(this);
        UnaryOps operator = unaryOp.getOperator();

        if (operator == UnaryOps.LNEG) {
            if (type != TypeTable.getType("boolean")) {
                throw new SemanticError(unaryOp.getLine(), unaryOp.getOperator().toString() + " with a non-boolean operand");
            }
            return TypeTable.getType("boolean");
        } else {
            if (type != TypeTable.getType("int"))
                throw new SemanticError(unaryOp.getLine(), "Unary " + unaryOp.getOperator().toString() + " with a non-int operand");
            return TypeTable.getType("int");
        }

    }

    @Override
    public Object visit(GlobalSymbolTable table) {
        return null;
    }

    @Override
    public Object visit(ClassSymbolTable table) {
        return null;
    }

    @Override
    public Object visit(BlockSymbolTable table) {
        return null;
    }

    @Override
    public Object visit(MethodSymbolTable table) {
        return null;
    }

    @Override
    public Object visit(SymbolTable symbolTable) {
        return null;
    }

}
